"""
This code is based on the tutorial on deeplearning.org


References:

    - textbooks: "Pattern Recognition and Machine Learning" -
                 Christopher M. Bishop, section 4.3.2

"""
__docformat__ = 'restructedtext en'

import cPickle
import gzip
import os
import sys
import time

import numpy
import nnutils

import theano
import theano.tensor as T


class LogisticRegression(object):
    """Multi-class Logistic Regression Class

    The logistic regression is fully described by a weight matrix :math:`W`
    and bias vector :math:`b`. Classification is done by projecting data
    points onto a set of hyperplanes, the distance to which is used to
    determine a class membership probability.
    """

    def __init__(self, input, input_dropout,
                 n_in, n_out, numpy_rng, theano_rng, 
                 W=None,b=None,L1rec=None,L2rec=None,
                 dropout_rate = None):
        """ Initialize the parameters of the logistic regression

        :type input: theano.tensor.TensorType
        :param input: symbolic variable that describes the input of the
                      architecture (one minibatch)

        :type n_in: int
        :param n_in: number of input units, the dimension of the space in
                     which the datapoints lie

        :type n_out: int
        :param n_out: number of output units, the dimension of the space in
                      which the labels lie

        """

        # initialize with 0 the weights W as a matrix of shape (n_in, n_out)
        if W == None:
             self.W = theano.shared(value=numpy.zeros((n_in, n_out),
                                                  dtype=theano.config.floatX),
                                 name='W', borrow=True)
        else:
            self.W = W
        # initialize the baises b as a vector of n_out 0s
        if b == None:
            self.b = theano.shared(value=numpy.zeros((n_out,),
                                                 dtype=theano.config.floatX),
                               name='b', borrow=True)
        else:
            self.b = b
        
        self.dropout_rate = dropout_rate
        # compute vector of class-membership probabilities in symbolic form
        if dropout_rate == None:
            self.p_y_given_x = T.nnet.softmax(T.dot(input, self.W) + self.b)
            # compute prediction as class whose probability is maximal in
            # symbolic form
            self.y_pred = T.argmax(self.p_y_given_x, axis=1)
        else:
            dropout_h = T.dot(input_dropout \
                              * T.cast(
                                   theano_rng.binomial(
                                        size=input_dropout.shape, 
                                        n=1, p=1-dropout_rate
                                        ), 
                                   dtype=theano.config.floatX
                                   ),  \
                              self.W)
            h = T.dot(input, self.W * (1-dropout_rate))
            self.p_y_given_x_dropout = T.nnet.softmax(dropout_h + self.b)
            self.p_y_given_x = T.nnet.softmax(h + self.b)
            self.y_pred = T.argmax(self.p_y_given_x, axis=1)

        # parameters of the model
        self.params = [self.W, self.b]
        
        if L1rec == None:
            self.L1rec = abs(self.W).sum()
        else:
            self.L1rec = L1rec + abs(self.W).sum()
            
        
        if L2rec == None:
            self.L2rec = (self.W **2).sum()
        else:
            self.L2rec = L2rec + (self.W **2).sum()

    def cost(self, y,L1_reg,L2_reg):
        """Return the mean of the negative log-likelihood of the prediction
        of this model under a given target distribution.

        .. math::

            \frac{1}{|\mathcal{D}|} \mathcal{L} (\theta=\{W,b\}, \mathcal{D}) =
            \frac{1}{|\mathcal{D}|} \sum_{i=0}^{|\mathcal{D}|} \log(P(Y=y^{(i)}|x^{(i)}, W,b)) \\
                \ell (\theta=\{W,b\}, \mathcal{D})

        :type y: theano.tensor.TensorType
        :param y: corresponds to a vector that gives for each example the
                  correct label

        Note: we use the mean instead of the sum so that
              the learning rate is less dependent on the batch size
        """
        # y.shape[0] is (symbolically) the number of rows in y, i.e.,
        # number of examples (call it n) in the minibatch
        # T.arange(y.shape[0]) is a symbolic vector which will contain
        # [0,1,2,... n-1] T.log(self.p_y_given_x) is a matrix of
        # Log-Probabilities (call it LP) with one row per example and
        # one column per class LP[T.arange(y.shape[0]),y] is a vector
        # v containing [LP[0,y[0]], LP[1,y[1]], LP[2,y[2]], ...,
        # LP[n-1,y[n-1]]] and T.mean(LP[T.arange(y.shape[0]),y]) is
        # the mean (across minibatch examples) of the elements in v,
        # i.e., the mean log-likelihood across the minibatch.
        if self.dropout_rate == None:
            c = -T.mean(T.log(self.p_y_given_x)[T.arange(y.shape[0]), y]) \
                + L1_reg*self.L1rec + L2_reg*self.L2rec
        else:
            c = -T.mean(T.log(self.p_y_given_x_dropout)[T.arange(y.shape[0]),
                y]) + L1_reg*self.L1rec + L2_reg*self.L2rec
            
        return c

    def errors(self, y):
        """Return a float representing the number of errors in the minibatch
        over the total number of examples of the minibatch ; zero one
        loss over the size of the minibatch

        :type y: theano.tensor.TensorType
        :param y: corresponds to a vector that gives for each example the
                  correct label
        """

        # check if y has same dimension of y_pred
        if y.ndim != self.y_pred.ndim:
            raise TypeError('y should have the same shape as self.y_pred',
                ('y', y.type, 'y_pred', self.y_pred.type))
        # check if y is of the correct datatype
        if y.dtype.startswith('int'):
            # the T.neq operator returns a vector of 0s and 1s, where 1
            # represents a mistake in prediction
            return T.mean(T.neq(self.y_pred, y))
        else:
            raise NotImplementedError()

